package com.example.demo.controller;

import dev.langchain4j.model.chat.StreamingChatLanguageModel;
import dev.langchain4j.model.chat.response.ChatResponse;
import dev.langchain4j.model.chat.response.StreamingChatResponseHandler;
import dev.langchain4j.model.openai.OpenAiStreamingChatModel;
import jakarta.validation.constraints.NotBlank;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.web.bind.annotation.*;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Sinks;

import java.util.concurrent.atomic.AtomicBoolean;

@RestController
@RequestMapping("/snow-chat-ai")
public class AIController {

    private final static Logger logger = LoggerFactory.getLogger(AIController.class);

    private final StreamingChatLanguageModel streamingChatLanguageModel;

    private final Sinks.Many<String> sink = Sinks.many().unicast().onBackpressureBuffer();

    public AIController(@Value("${langchain4j.open-ai.chat-model.api-key}") String aiKey,
                        StreamingChatLanguageModel streamingChatLanguageModel) {
        this.streamingChatLanguageModel = OpenAiStreamingChatModel.builder()
                .apiKey(System.getenv(aiKey)).modelName("qwen-plus")
                .build();
    }

    @PostMapping("/user-chat")
    public Flux<String> chat(@RequestParam @NotBlank String message) {
        System.out.println(message);
        AtomicBoolean isCompleted = new AtomicBoolean(false);
        Sinks.Many<String> requestSink = Sinks.many().unicast().onBackpressureBuffer();

        streamingChatLanguageModel.chat(message, new StreamingChatResponseHandler() {
            @Override
            public void onPartialResponse(String partialResponse) {
                if (!isCompleted.get()) {
                    requestSink.tryEmitNext(partialResponse);
                }
            }

            @Override
            public void onCompleteResponse(ChatResponse completeResponse) {
                if (!isCompleted.get()) {
                    requestSink.tryEmitNext(completeResponse.aiMessage().text());
                    isCompleted.set(true);
                    requestSink.tryEmitComplete();
                }
            }

            @Override
            public void onError(Throwable error) {
                if (!isCompleted.get()) {
                    String errorMessage = error.getMessage();
                    Sinks.EmitResult result = requestSink.tryEmitNext(errorMessage);
                    if (result.isFailure()) {
                        logger.warn("Failed to emit partial response: {}", result);
                    }
                    isCompleted.set(true);
                    requestSink.tryEmitComplete();
                }
            }
        });

        return requestSink.asFlux();
    }

}
